import gdal, osr
import numpy as np
import os
from src.utils.loghelper import logHelper
from src.utils.utils import find_layers_files, merge_modis_to_tiff, round_up


class AGRO5DownscalingPoint5Degree(object):

    _x_final = 720
    _y_final = 360

    def process(self, file_path, out_put_file_path):

        files = os.listdir(file_path)

        for file in files:
            result_array_1 = np.zeros((self._y_final, self._x_final))
            result_array_2 = np.zeros((self._y_final, self._x_final))
            result_array_3 = np.zeros((self._y_final, self._x_final))

            data_set = gdal.Open(file_path + "/" + file)

            geo_transform = data_set.GetGeoTransform()

            x_size = data_set.RasterXSize
            y_size = data_set.RasterYSize

            x_start = geo_transform[0]
            x_end = geo_transform[0] + x_size * geo_transform[1]
            y_start = geo_transform[3] + geo_transform[5] * y_size
            y_end = geo_transform[3]
            # y_origin = geo_transform[3] - y_size * geo_transform[5]

            data_array = data_set.ReadAsArray(0, 0, x_size, y_size)

            x_extend = x_end - x_start
            y_extend = y_end - y_start

            x_count = round(x_size/(x_extend/0.5))
            y_count = round(y_size/(y_extend/0.5))

            for i in range(0, 360):
                x_extend_start = i * x_count
                x_extend_end = (i + 1) * x_count

                for j in range(0, 720):

                    y_extend_start = j * y_count
                    y_extend_end = (j + 1) * y_count

                    if x_extend_end > x_size:
                        x_extend_end = x_size
                    if y_extend_end > y_size:
                        y_extend_end = y_size

                    if y_extend_start == y_extend_end:
                        break
                    else:

                        input_array = data_array[x_extend_start: x_extend_end, y_extend_start: y_extend_end]
                        if 0 not in input_array.shape:
                            result_array_1[i, j], \
                            result_array_2[i, j], \
                            result_array_3[i, j] = self._calculate(
                                data_array[x_extend_start: x_extend_end, y_extend_start: y_extend_end])

                if x_extend_start == x_extend_end:
                    break

            driver = gdal.GetDriverByName('GTiff')

            out_raster = driver.Create(
                out_put_file_path + "/" + file.split(".")[0] + "_5downscaling.tif", self._x_final, self._y_final, 3, gdal.GDT_Float32)

            logHelper.info(out_put_file_path + "/" + file.split(".")[0] + "_5downscaling.tif" + "has been created")

            target_xsize = 0.5
            target_ysize = 0.5
            out_raster.SetGeoTransform((-180, target_xsize, 0, 90, 0, target_ysize))

            # spatial reference
            outRasterSRS = osr.SpatialReference()
            # 代码4326表示WGS84坐标
            outRasterSRS.ImportFromEPSG(4326)
            out_raster.SetProjection(outRasterSRS.ExportToWkt())

            # for i in range(1, 4):
            out_band = out_raster.GetRasterBand(1)
            out_band.WriteArray(result_array_1)
            out_band = out_raster.GetRasterBand(2)
            out_band.WriteArray(result_array_2)
            out_band = out_raster.GetRasterBand(3)
            out_band.WriteArray(result_array_3)

            out_raster.FlushCache()

    def _calculate(self, array):

        [rows, cols] = array.shape
        raster_size = rows * cols

        array_1 = np.where(array == 1)
        array_2 = np.where(array == 2)
        array_3 = np.where(array == 3)

        return len(array_1[0]) / raster_size, len(array_2[0]) / raster_size, len(array_3[0]) / raster_size


class AGRO5Downscaling1Degree(object):

    _x_final = 720
    _y_final = 360

    _file_name = None
    _target_file_name = "5downscaling" + ".tif"
    _i = 1

    def process(self, file_path):

        # initial result
        result_array_1 = np.zeros((self._y_final, self._x_final))
        result_array_2 = np.zeros((self._y_final, self._x_final))
        result_array_3 = np.zeros((self._y_final, self._x_final))

        # read file

        files = os.listdir(file_path)

        for file in files:
            print("%d  %d" % (self._i, len(files)))

            self._file_name = file.split(".")[0]

            data_set = gdal.Open(file_path + "\\" + file)

            rows = data_set.RasterXSize
            columns = data_set.RasterYSize

            mat = data_set.GetGeoTransform()
            x_origin = round(mat[0])
            y_origin = round(mat[3] - 2)

            data_array = data_set.ReadAsArray(0, 0, data_set.RasterYSize, data_set.RasterYSize)

            # create result array
            temp_array_1 = np.zeros((4, 4))
            temp_array_2 = np.zeros((4, 4))
            temp_array_3 = np.zeros((4, 4))

            for i in range(0, 4):
                for j in range(0, 4):

                    if i < 3:
                        x_extend_start = i * 1800
                        x_extend_end = x_extend_start + 1800
                    else:
                        x_extend_start = i * 1800 + 1
                        x_extend_end = x_extend_start + 1799

                    if j < 3:
                        y_extend_start = j * 1800
                        y_extend_end = y_extend_start + 1800
                    else:
                        y_extend_start = j * 1800 + 1
                        y_extend_end = y_extend_start + 1799

                    temp_array_1[i, j], \
                    temp_array_2[i, j], \
                    temp_array_3[i, j] = self._calculate(
                        data_array[x_extend_start: x_extend_end, y_extend_start: y_extend_end])

                    y_offset = y_origin + 0.5 * j
                    x_offset = x_origin + 0.5 * i

                    result_array_1[int((90 + y_offset) * 2), int((180 + x_offset) * 2)] = temp_array_1[i, j]
                    result_array_2[int((90 + y_offset) * 2), int((180 + x_offset) * 2)] = temp_array_2[i, j]
                    result_array_3[int((90 + y_offset) * 2), int((180 + x_offset) * 2)] = temp_array_3[i, j]

            self._i = self._i + 1

        driver = gdal.GetDriverByName('GTiff')
        out_raster = driver.Create(
            "D:\\AGROTestData" + "\\" + self._target_file_name, self._x_final, self._y_final, 3,
            gdal.GDT_Float32)

        target_xsize = 0.5
        target_ysize = 0.5
        out_raster.SetGeoTransform((-180, target_xsize, 0, 90, 0, target_ysize))

        # spatial reference
        outRasterSRS = osr.SpatialReference()
        # 代码4326表示WGS84坐标
        outRasterSRS.ImportFromEPSG(4326)
        out_raster.SetProjection(outRasterSRS.ExportToWkt())

        # for i in range(1, 4):
        out_band = out_raster.GetRasterBand(1)
        out_band.WriteArray(result_array_1)
        out_band = out_raster.GetRasterBand(2)
        out_band.WriteArray(result_array_2)
        out_band = out_raster.GetRasterBand(3)
        out_band.WriteArray(result_array_3)

        out_raster.FlushCache()

    def _calculate(self, array):

        [rows, cols] = array.shape
        raster_size = rows * cols

        array_1 = np.where(array == 1)
        array_2 = np.where(array == 2)
        array_3 = np.where(array == 3)

        return len(array_1[0]) / raster_size, len(array_2[0]) / raster_size, len(array_3[0]) / raster_size


class AGRO5DownscalingForLandsatPoint5Degree(object):

    _x_final = 720
    _y_final = 360

    _file_name = None
    _target_file_name = "5downscaling" + ".tif"
    _i = 1

    def process(self, file_path):

        # initial result
        result_array_1 = np.zeros((self._y_final, self._x_final))
        result_array_2 = np.zeros((self._y_final, self._x_final))
        result_array_3 = np.zeros((self._y_final, self._x_final))

        # read file

        files = os.listdir(file_path)

        for file in files:

            print("%d  %d" %(self._i, len(files)))

            self._file_name = file.split(".")[0]

            data_set = gdal.Open(file_path + "\\" + file)

            mat = data_set.GetGeoTransform()
            x_origin = mat[0]
            y_origin = mat[3]

            data_array = data_set.ReadAsArray(0, 0, data_set.RasterYSize, data_set.RasterYSize)

            # create result array
            temp_array_1 = np.zeros((2, 2))
            temp_array_2 = np.zeros((2, 2))
            temp_array_3 = np.zeros((2, 2))

            for i in range(0, 2):
                for j in range(0, 2):

                    if i == 0:
                        x_extend_start = i * 1800
                        x_extend_end = x_extend_start + 1800
                    else:
                        x_extend_start = i * 1800 + 1
                        x_extend_end = x_extend_start + 1799

                    if j == 0:
                        y_extend_start = j * 1800
                        y_extend_end = y_extend_start + 1800
                    else:
                        y_extend_start = j * 1800 + 1
                        y_extend_end = y_extend_start + 1799

                    temp_array_1[i, j], \
                    temp_array_2[i, j], \
                    temp_array_3[i, j] = self._calculate(
                        data_array[x_extend_start: x_extend_end, y_extend_start: y_extend_end])

                    y_offset = y_origin + 0.5 * j
                    x_offset = x_origin + 0.5 * i

                    result_array_1[int(round_up((90 - y_offset) * 2)), int(round_up((180 + x_offset) * 2))] = temp_array_1[i, j]
                    result_array_2[int(round_up((90 - y_offset) * 2)), int(round_up((180 + x_offset) * 2))] = temp_array_2[i, j]
                    result_array_3[int(round_up((90 - y_offset) * 2)), int(round_up((180 + x_offset) * 2))] = temp_array_3[i, j]

            self._i = self._i + 1

        driver = gdal.GetDriverByName('GTiff')
        out_raster = driver.Create(
            "D:\\AGROTestData" + "\\" + self._target_file_name, self._x_final, self._y_final, 1,
            gdal.GDT_Float32)

        target_xsize = 0.5
        target_ysize = 0.5
        out_raster.SetGeoTransform((-180, target_xsize, 0, 90, 0, target_ysize))

        # spatial reference
        outRasterSRS = osr.SpatialReference()
        # 代码4326表示WGS84坐标
        outRasterSRS.ImportFromEPSG(4326)
        out_raster.SetProjection(outRasterSRS.ExportToWkt())

        # for i in range(1, 4):
        out_band = out_raster.GetRasterBand(1)
        out_band.WriteArray(result_array_1)
        # out_band = out_raster.GetRasterBand(2)
        # out_band.WriteArray(result_array_2)
        # out_band = out_raster.GetRasterBand(3)
        # out_band.WriteArray(result_array_3)

        out_raster.FlushCache()

    def merge_files(self):

        # merged_output_file_path = os.getcwd() + "\\tempGlobal" + "\\mergedTiff"
        # merged_output_path = os.getcwd() + "\\tempGlobal" + "\\mergedTiff" + "\\global.tif"

        merged_output_file_path = "D:\\AGROTestData\\temp"
        merged_output_path = merged_output_file_path + "\\global.tif"

        # tiff_list = find_layers_files(os.getcwd() + "\\temp1", '.tif')
        tiff_list = find_layers_files("D:\AGROTestData\\2010", '.tif')
        if os.path.exists(merged_output_file_path):
            if os.path.exists(merged_output_path):
                os.remove(merged_output_path)
                print('{} tiff file is exist and remove'.format(merged_output_path))
        else:
            os.makedirs(merged_output_file_path)
        merge_modis_to_tiff(tiff_list, merged_output_path, 0)

    def _calculate(self, array):

        [rows, cols] = array.shape
        raster_size = rows * cols

        array_1 = np.where(array == 1)
        array_2 = np.where(array == 2)
        array_3 = np.where(array == 3)

        return len(array_1[0]) / raster_size, len(array_2[0]) / raster_size, len(array_3[0]) / raster_size



